package main

import (
	"crypto/ecdh"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"log"
	"net/http"
	"os"
	"bytes"
	"io"
	"net"
	// "encoding/base64"

	"golang.org/x/net/dns/dnsmessage"
)

func main() {

	// 读取并解析ECDH私钥
	pemData, err := os.ReadFile("ecdh_private_key.pem")
	if err != nil {
		fmt.Println("Error reading PEM file:", err)
		return
	}
	block, _ := pem.Decode(pemData)
	if block == nil {
		fmt.Println("Error decoding PEM block")
		return
	}
	ecdhSKBytes, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		fmt.Println("failed to marshal private key into PKIX format")
		return
	}
	privateKey := ecdhSKBytes.(*ecdh.PrivateKey)

	// 获取ECH配置列表
	// echConfigList, err := GetECHConfigList(privateKey, []string{"server0.com"})
	// if err != nil {
	// 	fmt.Println("failed to get echconfiglist")
	// 	return
	// }
	// fmt.Println(echConfigList)

	// 获取特定域名的ECH配置信息
	echConfig, err := GetECHConfig(privateKey, "server0.com")
	if err != nil {
		fmt.Println("failed to get echConfig")
		return
	}

	// fmt.Printf("ECHConfig: %s\n", base64.StdEncoding.EncodeToString(echConfig))
	// fmt.Printf("echConfigList Std: %s\n", base64.StdEncoding.EncodeToString(echConfigList))

	// 配置TLS
	tlsConfig := &tls.Config{
		GetCertificate:   getCertificateForSNI,
		MinVersion:       tls.VersionTLS13,
		MaxVersion:       tls.VersionTLS13,
		CurvePreferences: []tls.CurveID{},
		EncryptedClientHelloKeys: []tls.EncryptedClientHelloKey{{
			Config:      echConfig,
			PrivateKey:  privateKey.Bytes(),
			SendAsRetry: true,
		}},
	}

	// 创建一个http.server实例
	server := &http.Server{
		Addr:      ":443",
		TLSConfig: tlsConfig,
	}

	// 注册 DoH 请求处理函数
	http.HandleFunc("/dns-query", handleECHRequest)

	// 启动 HTTPS 服务器
	log.Println("Starting DoH server on https://localhost:443")
	if err := server.ListenAndServeTLS("", ""); err != nil {
		log.Fatal("ListenAndServeTLS: ", err)
	}

}

func handleECHRequest(w http.ResponseWriter, r *http.Request) {
	// 读取 DoH+Ech 请求中的 DNS 消息
	dnsRequest := r.Body
	defer dnsRequest.Close()

	// 读取请求中的 DNS 消息内容
	var buf bytes.Buffer
	_, err := io.Copy(&buf, dnsRequest)
	if err != nil {
		log.Println("Failed to read DNS request: ", err)
		http.Error(w, "Failed to read DNS request", http.StatusInternalServerError)
		return
	}

	// 解包 DNS 消息
	var msg dnsmessage.Message
	err = msg.Unpack(buf.Bytes())
	if err != nil {
		log.Println("Failed to unpack DNS request: ", err)
		http.Error(w, "Failed to unpack DNS request", http.StatusInternalServerError)
		return
	}

	// 向真实的 DNS 服务器发送请求
	dnsResponse, err := sendDNSRequestToRealServer(msg)
	if err != nil {
		log.Println("Failed to send DNS request to real server: ", err)
		http.Error(w, "Failed to get DNS response", http.StatusInternalServerError)
		return
	}

	// 发送 DNS 响应作为 DoH 响应
	w.Header().Set("Content-Type", "application/dns-message")
	w.Write(dnsResponse)
}

// 向真实的 DNS 服务器发送 DNS 请求
func sendDNSRequestToRealServer(msg dnsmessage.Message) ([]byte, error) {
	// 连接到 DNS 服务器
	dnsServer := "114.114.114.114:53"
	conn, err := net.Dial("udp", dnsServer)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to DNS server: %w", err)
	}
	defer conn.Close()

	// 将 DNS 消息序列化
	dnsBytes, err := msg.Pack()
	if err != nil {
		return nil, fmt.Errorf("failed to pack DNS message: %w", err)
	}

	// 发送 DNS 请求
	_, err = conn.Write(dnsBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to send DNS request: %w", err)
	}

	// 读取 DNS 响应
	responseBuf := make([]byte, 512)
	n, err := conn.Read(responseBuf)
	if err != nil {
		return nil, fmt.Errorf("failed to read DNS response: %w", err)
	}

	return responseBuf[:n], nil
}

func getCertificateForSNI(helloInfo *tls.ClientHelloInfo) (*tls.Certificate, error) {
	switch helloInfo.ServerName {
	case "server1.com":
		cert, err := tls.LoadX509KeyPair("server1.crt", "server1.key")
		return &cert, err // 返回结构体的指针
	case "server2.com":
		cert, err := tls.LoadX509KeyPair("server2.crt", "server2.key")
		return &cert, err // 返回结构体的指针
	default:
		return nil, fmt.Errorf("no certificate found for SNI: %s", helloInfo.ServerName)
	}
}